#include "utils.h"
#include "bessel.c"
#include "casts.cpp"

const double PI = 3.14159265358979323846264338327950288;
const double HRED = 1.;

int prnt;

// grid properties
int n, m, n2, n3;
double dT, dX, dP, dE;
double minP, maxP, M;
double Vmax, Vmin;
double dimension;        

// Kosloff absorption
double absorptionDecay;
double absorptionPotential;

// bessel coeffs
double *bjk;

// radiation fied
double photonEnergy;
double electricFieldInt;
double dipoleMoment;
double impulseArea;
double impulsePeriod;

// phase shift
fftw_complex Phase;

// FFTW
fftw_complex* out;
fftw_plan p, q;

//
fftw_complex *suma_czebT, *czeb_n0T, *czeb_n1T, *czeb_n2T, *czeb_n1_hamT;

fftw_complex* HPsiT;
fftw_complex* Vcalc;
Point* CoordT;

// wavefunction evolution output file
FILE* evolution;

void printPsi(fftw_complex *PsiT, double t)
{
	FILE * pFile;
	char fname[50];
	
	sprintf(fname, "%scart.%d", ctlg, (int)(1000.*t+0.99));
	//	string name = "3DHAbs10\\cart.";
	//	name = name+itoa((int)(1000.*t+0.99));
	
	printf(" - opening data output file: %s", fname);
	//	pFile = fopen(name.c_str(),"w");
	pFile = fopen(fname,"w");
	
	fprintf(pFile,"#Wavefunction at %.7f\n----------------------------\n#norm: %.4f\n----------------------------\n", t, calcNorm(PsiT, n3, dX));
	
	FOR(i,n2)
	{
		fprintf(pFile, "%.5f %.5f %.5f %e %e %e\n", CoordT[i].x, CoordT[i].y, CoordT[i].z, PsiT[i][0], PsiT[i][1], sqrt(PsiT[i][0]*PsiT[i][0]+PsiT[i][1]*PsiT[i][1]));
	}
	
	fprintf(pFile, "\n#==================\n");
	
	fclose (pFile);
}

void printEvolution(fftw_complex* Psi, fftw_complex* PsiInit, fftw_complex* PsiFinal1, fftw_complex* PsiFinal2, fftw_complex* PsiFinal3, double t)
{
	if(evolution == NULL) {
		char fname[50];
		sprintf(fname, "%sevol", ctlg);
		evolution = fopen(fname, "w");
	}
	
	fftw_complex castT; calcCast(PsiInit, Psi, n3, dX, castT);
	fftw_complex cast1T; calcCast(PsiFinal1, Psi, n3, dX, cast1T);
	fftw_complex cast2T; calcCast(PsiFinal2, Psi, n3, dX, cast2T);
	fftw_complex cast3T; calcCast(PsiFinal3, Psi, n3, dX, cast3T);
	/*
	fcomplex cast1 = calcCast(PsiFinal1, Psi, dX);
	fcomplex cast2 = calcCast(PsiFinal2, Psi, dX);
	fcomplex cast3 = calcCast(PsiFinal3, Psi, dX);*/
	fprintf(evolution, "%.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.7f\n", t, castT[0], castT[1], cast1T[0], cast1T[1], cast2T[0], cast2T[1], cast3T[0], cast3T[1], calcNorm(Psi, n3, dX));
	// fprintf(evolution, "%.4f %.4f %.4f %.7f\n", t, castT[0], castT[1], calcNorm(Psi, n3, dX));
}

void assignHydrogenAtom(int Qn, int Ql, int Qm, fftw_complex* PsiT) {
	double r, x, y, z;
	if(Qn==1 && Ql==0 && Qm==0) {
		double Nrm = 0.564189584;
		FOR(i,n3) {
			x = CoordT[i].x, y = CoordT[i].y, z = CoordT[i].z; 
			r = sqrt(x*x+y*y+z*z);
			PsiT[i][0] = Nrm*exp(-r);
			PsiT[i][1] = 0.;
		}
	}
	else if(Qn==2 && Ql==1 && Qm==0) {
		double Nrm = 0.0997355701;
		FOR(i,n3) {
			x = CoordT[i].x, y = CoordT[i].y, z = CoordT[i].z; 
			r = sqrt(x*x+y*y+z*z);
			PsiT[i][0] = Nrm*z*exp(-r/2.);
			PsiT[i][1] = 0.;
		}
	}
	else if(Qn==3 && Ql==0 && Qm==0) {
		double Nrm = 0.00402141985;
		FOR(i,n3) {
			x = CoordT[i].x, y = CoordT[i].y, z = CoordT[i].z; 
			r = sqrt(x*x+y*y+z*z);
			PsiT[i][0] = Nrm*(27.-18.*r+2.*r*r)*exp(-r/3.);
			PsiT[i][1] = 0.;
		}
	}
	else if(Qn==3 && Ql==1 && Qm==0) {
		double Nrm =0.00985042668;
		FOR(i,n3) {
			x = CoordT[i].x, y = CoordT[i].y, z = CoordT[i].z; 
			r = sqrt(x*x+y*y+z*z);
			PsiT[i][0] = Nrm*(6.-r)*z*exp(-r/3.);
			PsiT[i][1] = 0.;
		}
	}
	else if(Qn==3 && Ql==2 && Qm==0) {
		double Nrm = 0.00284357325;
		FOR(i,n3) {
			x = CoordT[i].x, y = CoordT[i].y, z = CoordT[i].z; 
			r = sqrt(x*x+y*y+z*z);
			PsiT[i][0] = Nrm*(3.*z*z-r*r)*exp(-r/3.);
			PsiT[i][1] = 0.;
		}
	}
}

//      Hydrogen-ish potential
void V(double t, const Point &p, fftw_complex &q) {
	double r=sqrt(p.x*p.x+p.y*p.y+p.z*p.z);
	double pulse = 1.;
	
	if(t>impulsePeriod)
	pulse = 0.;    
	
	// to remove the external electric field set electricField=0.
	if(r>1.e-2) {
		//		q[0] = -1./r+pulse*electricFieldInt*dipoleMoment*cos(photonEnergy*t), q[1] = 0.;
		q[0] = -1./r+pulse*electricFieldInt*p.z*cos(photonEnergy*t), q[1] = 0.;
	}
	else {
		q[0] = -1.e2+pulse*electricFieldInt*p.z*cos(photonEnergy*t), q[1] = 0.;
	}
}

void calcDE(double t)
{
	/*
	Vmax=norm(V(t, Point(0.,0.,0.))), Vmin=Vmax;
	double v;
	
	for(VVVP::iterator i=Coord.begin(); i!=Coord.end(); i++)
	for(VVP::iterator j=i->begin(); j!=i->end(); j++)
	for(VP::iterator k=j->begin(); k!=j->end(); k++) {
		v=norm(V(t, *k));
		if(v<Vmin) Vmin=v;
		else if(v>Vmax) Vmax=v;
	}
	*/
	
	dE=electricFieldInt*dipoleMoment*2+maxP*maxP/M+1.e2;
	
	Vmin = -1.e-2-electricFieldInt*dipoleMoment;
	double s=(0.5 * dE + Vmin) * dT / HRED;
	Phase[0] = cos(s), Phase[1] = -sin(s);
	
	char fname[50];
	sprintf(fname, "%sbessel", ctlg);
	
	FILE* bessel = fopen(fname,"w");
	fprintf(bessel, "# Bessel(%.6f)\n",dE*dT/(2.*HRED));    
	FOR(i,m) {
		bjk[i]=BesselJ(i,dE*dT/(2.*HRED))*(2.-1.*(i==0));
		fprintf(bessel, "%d %.8f\n",i,bjk[i]);  
	}
	fclose(bessel);
}

// FFTW functions
void kineticEnergyOperator(fftw_complex* Psi)
{	
	if(out=='\0') {
		out = new fftw_complex[n3];
		p = fftw_plan_dft_3d(n, n, n, Psi, out, FFTW_FORWARD, FFTW_ESTIMATE);
		q = fftw_plan_dft_3d(n, n, n, out, Psi, FFTW_BACKWARD, FFTW_ESTIMATE);
	}
	
	fftw_execute(p);
	
	double px, py, pz, mlt = 1./(2.*M*n3);
	FOR(i,n3) {
		px = (2*(i/n2)<n? (i/n2) : ((i/n2)-n))*dP;
		py = (2*((i/n)%n)<n? ((i/n)%n) : ((i/n)%n-n))*dP;
		pz = (2*(i%n)<n? (i%n) : (i%n-n))*dP;
		out[i][0] = out[i][0]*(px*px+py*py+pz*pz)*mlt;
		out[i][1] = out[i][1]*(px*px+py*py+pz*pz)*mlt;
	}
	fftw_execute(q);
	
//	fftw_destroy_plan(p);
//	fftw_destroy_plan(q);
	//	fftw_free(out);
}

void Hamiltonian(fftw_complex* Psi, fftw_complex* HPsi, double t) {
	
	kineticEnergyOperator(HPsi);
	
	FOR(i,n3) {
		HPsi[i][0]=HPsi[i][0]+Vcalc[i][0]*Psi[i][0]-Vcalc[i][1]*Psi[i][1];
		HPsi[i][1]=HPsi[i][1]+Vcalc[i][0]*Psi[i][1]+Vcalc[i][1]*Psi[i][0];
	}
}

// Kosloff absorption
void applyAbsorption(fftw_complex* Psi) {
	double abs;
	double s, x, y, z;
	FOR(i,n3)
	{
		x = CoordT[i].x, y = CoordT[i].y, z = CoordT[i].z; 
		s = sqrt(x*x+y*y+z*z);		
		s=(double)n/2-s;
		if(s*absorptionDecay>5.) {
			continue;
		}
		abs = cosh(s*absorptionDecay);
		abs = dT * absorptionPotential / (abs*abs);
		Psi[i][0] = Psi[i][0] * (1.-abs);
		Psi[i][1] = Psi[i][1] * (1.-abs);
	}
}

void singleStep(fftw_complex* czeb_n0T, double t)
{
	printf("-----------------------------------------\nStep: %.4f\n",t);
	
	FOR(i,n3)
	V(t, CoordT[i], Vcalc[i]);
	
	if(czeb_n1T == NULL) {
		czeb_n1T =	(fftw_complex *)malloc(sizeof(fftw_complex)*n3);
		suma_czebT =	(fftw_complex *)malloc(sizeof(fftw_complex)*n3);	
		czeb_n1_hamT =	(fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	}
	
	memcpy(czeb_n1T, czeb_n0T, n3*sizeof(fftw_complex));
	Hamiltonian(czeb_n0T, czeb_n1T, t);
	
	double s;
	double mlt = 2./dE, mlt2 = dE/2.+Vmin, mlt3 = -4./dE;
	FOR(i, n3) {
		s=czeb_n1T[i][0];
		czeb_n1T[i][0]=(czeb_n1T[i][1]-czeb_n0T[i][1]*mlt2)*mlt;
		czeb_n1T[i][1]=-(s-czeb_n0T[i][0]*mlt2)*mlt;
		suma_czebT[i][0]=czeb_n0T[i][0]*bjk[0]+czeb_n1T[i][0]*bjk[1];
		suma_czebT[i][1]=czeb_n0T[i][1]*bjk[0]+czeb_n1T[i][1]*bjk[1];
	}
	
	//rekurencja do obliczania Czeb
	int p = 0;
	fftw_complex *c0, *c1;
	FORA(j,2,m)
	{
		// printf("%d %d\n", j, p);
		c0 = (p == 0 ? czeb_n0T : czeb_n1T);
		c1 = (p == 0 ? czeb_n1T : czeb_n0T);
		memcpy(czeb_n1_hamT, c1, n3*sizeof(fftw_complex));
		Hamiltonian(c1,czeb_n1_hamT,t);				//H*Czeb_n1(-iH)*Psi
		FOR(i,n3) {
			c0[i][0] += mlt3 * (-czeb_n1_hamT[i][1]+c1[i][1]*mlt2);
			c0[i][1] += mlt3 * (czeb_n1_hamT[i][0]-c1[i][0]*mlt2);
			suma_czebT[i][0] += c0[i][0]*bjk[j];
			suma_czebT[i][1] += c0[i][1]*bjk[j];
		}
		p = (p==1 ? 0 : 1);		
	}	
	printf("\nCzebyshev calcs finished\n");
	
	FOR(i,n3) {
		czeb_n0T[i][0] = Phase[0]*suma_czebT[i][0] - Phase[1]*suma_czebT[i][1];
		czeb_n0T[i][1] = Phase[0]*suma_czebT[i][1] + Phase[1]*suma_czebT[i][0];
	}
	
	applyAbsorption(czeb_n0T);
	
	if((int)(1000.*t)%8000==0) {
		printf("Printing the results:\n");
		
		time_t pseconds1, pseconds2; time(&pseconds1);
		printPsi(czeb_n0T, t);
		time(&pseconds2);printf("Printing time: %.2f\n", difftime(pseconds2, pseconds1));
	}
}

void printEFieldData()
{
	char fname[50];
	sprintf(fname, "%semfield", ctlg);
	FILE* electric = fopen(fname, "w");      
	
	fprintf(electric, "Dipole approximation\n");
	fprintf(electric, "Transition dipole: %.4f\n", dipoleMoment);
	fprintf(electric, "----------------------------------\nRectangular impulse:\n");
	fprintf(electric, "Electric field: %.4f\n", electricFieldInt);
	fprintf(electric, "Ang. frequency: %.4f\t ; period: %.4f\n", photonEnergy, 2.*PI/photonEnergy);
	fprintf(electric, "Impulse width: %.4f\n", impulsePeriod);
	fprintf(electric, "----------------------------------\n");
	
	fclose(electric);
}


int main(int argc, char *argv[])
{
	time_t seconds1, seconds2;
	time(&seconds1);
	
	int NOSteps = 1;
	
	n=128;
	m=1;
	if(argc>1) m=atoi(argv[1]);
	if(argc>2) NOSteps=atoi(argv[2]);
	
	dimension = 200;
	
	evolution = NULL;
	
	dX = dimension/n;
	dT = 1.;
	dP = 2.*PI*HRED/(dX*(n-1));
	minP = -PI*HRED/dX;
	maxP = PI*HRED/dX;
	M = 1.;
	n3 = n*n*n;
	n2 = n*n;
	
	absorptionDecay = 1.;
	absorptionPotential = 1./dT;
	
	photonEnergy = 5./72.;
	electricFieldInt = 0.;//1.e-4;                       // 0.1-1. ; E||z
	dipoleMoment = 2.54689;         // dipole transition moment 320 <-> 210
	impulseArea = PI;
	
	Vcalc = (fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	CoordT = (Point *)malloc(sizeof(Point)*n3);
	
	if(electricFieldInt<1.e-5)
	impulseArea = 0.;
	impulsePeriod = impulseArea/(electricFieldInt*dipoleMoment);
	printEFieldData();
	
	fftw_complex* PsiT = (fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	fftw_complex* PsiInitT = (fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	fftw_complex* PsiFinal1T = (fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	fftw_complex* PsiFinal2T = (fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	fftw_complex* PsiFinal3T = (fftw_complex *)malloc(sizeof(fftw_complex)*n3);
	
	bjk = new double[m];
	FOR(i,n3) {
		CoordT[i] = Point(((double)((int)(i/n2)-n/2)+0.5)*dX,((double)((int)((i/n)%n)-n/2)+0.5)*dX,((double)((int)(i%n)-n/2)+0.5)*dX);
	}
	assignHydrogenAtom(2, 1, 0, PsiT);
	assignHydrogenAtom(2, 1, 0, PsiInitT);
	assignHydrogenAtom(3, 0, 0, PsiFinal1T);
	assignHydrogenAtom(3, 1, 0, PsiFinal2T);
	assignHydrogenAtom(3, 2, 0, PsiFinal3T);
	
	calcDE(0.);
	
	printf("dE: %e\n",dE);
	
	prnt=-1;
	
	printPsi(PsiT, 0.);
	printEvolution(PsiT, PsiInitT, PsiFinal1T, PsiFinal2T, PsiFinal3T, 0.);
	
	FOR(step, NOSteps) {
		singleStep(PsiT, dT*(step+1));
		printPsi(PsiT, dT*(step+1));
		printEvolution(PsiT, PsiInitT, PsiFinal1T, PsiFinal2T, PsiFinal3T, dT*(step+1));
	}
	free(out);
	
	fclose(evolution);
	
	time(&seconds2);
	
	printf("=====================\nTotal simulations time: %.2f s\n", difftime(seconds2, seconds1));
	
	return 0;
}
